import sklearn as sk
import numpy as np
import matplotlib
import pandas as pd
import matplotlib.pyplot as plt



from sklearn import feature_extraction

def one_hot_dataframe(data, cols, replace=False):
    """ Takes a dataframe and a list of columns that need to be encoded.
    Returns a 3-tuple comprising the data, the vectorized data,
    and the fitted vectorizor.
    Modified from https://gist.github.com/kljensen/5452382
    """
    vec = feature_extraction.DictVectorizer()
    mkdict = lambda row: dict((col, row[col]) for col in cols)
    
    #print 'Construyo vecData...'
    #print data[cols]
    #print cols

    # Create a dictionary for each row
    
    #print data[cols].apply(mkdict, axis=1).data
    #[0]['pclass']

    #vecData = pd.DataFrame(vec.fit_transform(data[cols].apply(mkdict, axis=1)).toarray())
    vecData = pd.DataFrame(vec.fit_transform(data[cols].to_dict(outtype='records')).toarray())
    vecData.columns = vec.get_feature_names()
    vecData.index = data.index
    if replace is True:
        data = data.drop(cols, axis=1)
        data = data.join(vecData)
    return (data, vecData)


def main():

    titanic = pd.read_csv('titanic.csv')
    print titanic

    raw_input("Press enter to continue")

    titanic, titanic_n= one_hot_dataframe(titanic, ['pclass', 'embarked', 'sex'], replace=True)

    titanic.describe()

if __name__ == "__main__":
    main()



